import networkx as nx  # 网络
from pylab import *
import pandas as pd
import os


# 获取数据分析结果，生成excel
def getresult(data, output_path):
    G = nx.from_pandas_edgelist(data, 'from', 'to', edge_attr='weight')
    # print(G)
    degree = nx.degree_centrality(G)
    closeness = nx.closeness_centrality(G)
    betweenness = nx.betweenness_centrality(G)
    eigenvector = nx.eigenvector_centrality(G)
    gd = nx.from_pandas_edgelist(data, 'from', 'to', edge_attr='weight', create_using=nx.DiGraph)
    indegree = gd.in_degree()
    outdegree = gd.out_degree()
    output = pd.DataFrame(
        columns=['country', 'degree', 'indegree', 'outdegree', 'closeness', 'betweenness', 'eigenvector'])
    output['country'] = G.nodes
    output['degree'] = [i for i in degree.values()]
    output['indegree'] = [i for i in indegree]
    output['outdegree'] = [i for i in outdegree]
    output['closeness'] = [i for i in closeness.values()]
    output['betweenness'] = [i for i in betweenness.values()]
    output['eigenvector'] = [i for i in eigenvector.values()]

    output.reset_index(inplace=True)
    excel_name = output_path + '.xlsx'
    # 保存数据分析excel表格
    output.to_excel(excel_name, index=False)
    # print(output)
    # 聚类系数
    nx.transitivity(G)
    nx.average_clustering(G)
    # 直径
    nx.diameter(G)
    # 所有节点间平均最短路径长度
    nx.average_shortest_path_length(G)
    print(excel_name, '图中所有的边数', G.nodes)
    print(excel_name, '图中边的条数', G.number_of_nodes())
    print(excel_name, '网络密度为：{}'.format(nx.density(G)))
    print(excel_name, '聚类系数为：{}'.format(nx.average_clustering(G)))
    print(excel_name, '直径为：：{}'.format(nx.diameter(G)))
    print(excel_name, '平均最短路径长度为：{}'.format(nx.average_shortest_path_length(G)))
    txt_name = output_path + "_print_info"
    txt_content = '图中所有的边数{}'.format(G.nodes)
    txt_content = txt_content + '\r\n' + '图中边的条数{}'.format(G.number_of_nodes())
    txt_content = txt_content + '\r\n' + '网络密度为：{}'.format(nx.density(G))
    txt_content = txt_content + '\r\n' + '聚类系数为：{}'.format(nx.average_clustering(G))
    txt_content = txt_content + '\r\n' + '直径为：：{}'.format(nx.diameter(G))
    txt_content = txt_content + '\r\n' + '平均最短路径长度为：{}'.format(nx.average_shortest_path_length(G))
    out_print_info(txt_name, txt_content)


# 将打印信息保存值txt文本中
def out_print_info(filename, content):
    # 获取当前文件所在路径
    file_os_path = os.getcwd()
    file_path = file_os_path + '\\' + filename + '.txt'
    file = open(file_path, 'w',  encoding='gb18030', errors='ignore')
    file.write(content)
    file.close()


# 轮询表中数据
def get_excel_data(excel_path):
    table_data = pd.read_excel(excel_path)
    # 读取到的excel表格打印
    # print(table_data)
    list_from = []
    list_to = []
    list_weight = []
    # 循环表格的每一行
    for i in range(0, table_data.shape[0]):
        # 循环列表的每一列
        for j in range(0, table_data.shape[1]):
            try:
                number = float(table_data.iloc[i, j])
                # 打印表中的每一个数字，从上往下，一行一行读取
                # print(number)
                # print("i = ", i, " j = ", j, "number = ", number)
                # print("list_title = ", list_title)
                # 如果当前元素大于0
                if number > 0:
                    # columns['from', 'to', 'weight']
                    # [i,j]究竟哪个为横坐标，哪个为纵坐标，先改了再说。
                    # 当前数字的列名
                    if excel_path.__contains__('6') or excel_path.__contains__('7'):
                        list_from.append(table_data.iloc[i, 1])
                    else:
                        list_from.append(table_data.iloc[i, 0])
                    # 当前数字的行名
                    list_to.append(table_data.columns[j])
                    # 当前数字的值
                    list_weight.append(table_data.iloc[i, j])
            except Exception as e:
                print("非数字格式 原数据：" + table_data.iloc[i, j])

    row = {"from": list_from, "to": list_to, "weight": list_weight}
    return pd.DataFrame(row)


# 根据文件路径获取excel表格数据分析 excel_path=excel表文件路径，output_path=分析后文件输出路径，filter_size=权重过滤值
def excel_analyse(excel_path, output_path, filter_size):
    excel_data = get_excel_data(excel_path)
    filter_data = excel_data[excel_data["weight"] > filter_size]
    print("筛选出weight > ", filter_size, "的数据")
    # 筛选后的数据打印
    print(filter_data)
    g = nx.from_pandas_edgelist(filter_data, 'from', 'to', edge_attr='weight', create_using=nx.DiGraph)
    edges = g.edges()
    weights = [g[u][v]['weight'] for u, v in edges]
    pos = nx.nx_pydot.graphviz_layout(g)
    pos_node_colors = get_node_colors(pos)
    plt.figure(figsize=(40, 40), dpi=400)
    nx.draw_networkx(g, pos, with_labels=True, width=weights, node_size=18000, font_size=50, edge_color='b',
                     linewidths=3, arrowsize=30, node_color=pos_node_colors, font_color='k', alpha=0.8)
    getresult(filter_data, output_path)

    plt.axis('off')
    file_path = output_path + '.png'
    plt.savefig(file_path, dpi=390)
    plt.show()


# 锁定不同国家颜色
def get_node_colors(pos):
    pos_keys = pos.keys()
    pos_node_colors = []
    for i in pos_keys:
        # 亚洲 紫色
        if i == "CHN" or i == "HKG" or i == "JPN" or i == "KOR" or i == "SGP" or i == "PHL" or i == "THA" \
                or i == "VNM" or i == "IND" or i == "TUR" or i == "MYS":
            node_color = "#673AB7"
            pos_node_colors.append(node_color)
        # 大洋洲 绿色
        elif i == "NZL" or i == "AUS":
            node_color = "#4CAF50"
            pos_node_colors.append(node_color)
        # 美洲 红色
        elif i == "USA" or i == "CAN" or i == "BRA" or i == "CHL" or i == "ARG" or i == "MEX" or i == "VEN":
            node_color = "#FF0000"
            pos_node_colors.append(node_color)
        # 欧洲 蓝色
        elif i == "RUS" or i == "ESP" or i == "NLD" or i == "FRA" or i == "DEU" or i == "SWE" or i == "CHE" \
                or i == "ITA" or i == "GBR" or i == "BEL" or i == "IRL" or i == "NOR" or i == "FIN" \
                or i == "PRT" or i == "LUX" or i == "GRC":
            node_color = "#2196F3"
            pos_node_colors.append(node_color)
        # 非洲 黄色
        elif i == "EGY" or i == "NGA" or i == "ARE":
            node_color = "#FFEB3B"
            pos_node_colors.append(node_color)
        # 中东地区 橙色
        elif i == "SAU" or i == "ISR":
            node_color = "#FF5722"
            pos_node_colors.append(node_color)
        else:
            print("未锁定区域的国家----", i)
            node_color = "#2196F3"
            pos_node_colors.append(node_color)
    return pos_node_colors


# 获取指定目录下的文件列表
def get_excel_file(dir_path):
    file_path_list = []
    for root, dirs, files in os.walk(dir_path):
        for file in files:
            if file.endswith("xlsx") or file.endswith("csv"):
                file_path_list.append(os.path.join(root, file))

    return file_path_list


if __name__ == '__main__':
    file_list = get_excel_file('file/')
    for i in file_list:
        if i.__contains__('5') or i.__contains__('6'):
            excel_analyse(i, i.replace('file', 'output'), 3)
        else:
            excel_analyse(i, i.replace('file', 'output'), 0.8)
